#include "SSB_lookup.h"
#include "SSB_sys_def.h"

#include <string.h>
#include <sys/param.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_system.h"
#include "esp_wifi.h"
#include "esp_event_loop.h"
#include "esp_log.h"
#include "nvs_flash.h"

#include "lwip/err.h"
#include "lwip/sockets.h"
#include "lwip/sys.h"
#include <lwip/netdb.h>

#define UDP_PORT 5000
#define MULTICAST_IPV4_ADDR ("224.0.0.251")

const int IPV4_GOTIP_BIT = BIT0;

static int socket_add_ipv4_multicast_group(int sock, bool assign_source_if)
{
    struct ip_mreq imreq = {0};
    struct in_addr iaddr = {0};
    int err = 0;
    imreq.imr_interface.s_addr = IPADDR_ANY;
    err = inet_aton(MULTICAST_IPV4_ADDR, &imreq.imr_multiaddr.s_addr);
    if (err != 1)
    {
        ESP_LOGE(SSB_TAG, "Configured IPV4 multicast address '%s' is invalid.", MULTICAST_IPV4_ADDR);
        goto err;
    }
    ESP_LOGI(SSB_TAG, "Configured IPV4 Multicast address %s", inet_ntoa(imreq.imr_multiaddr.s_addr));
    if (!IP_MULTICAST(ntohl(imreq.imr_multiaddr.s_addr)))
    {
        ESP_LOGW(SSB_TAG, "Configured IPV4 multicast address '%s' is not a valid multicast address. This will probably not work.", MULTICAST_IPV4_ADDR);
    }

    if (assign_source_if)
    {
        err = setsockopt(sock, IPPROTO_IP, IP_MULTICAST_IF, &iaddr,
                         sizeof(struct in_addr));
        if (err < 0)
        {
            ESP_LOGE(SSB_TAG, "Failed to set IP_MULTICAST_IF. Error %d", errno);
            goto err;
        }
    }

    err = setsockopt(sock, IPPROTO_IP, IP_ADD_MEMBERSHIP,
                     &imreq, sizeof(struct ip_mreq));
    if (err < 0)
    {
        ESP_LOGE(SSB_TAG, "Failed to set IP_ADD_MEMBERSHIP. Error %d", errno);
        goto err;
    }

err:
    return err;
}

static int create_multicast_ipv4_socket()
{
    struct sockaddr_in saddr = {0};
    int sock = -1;
    int err = 0;
    sock = socket(PF_INET, SOCK_DGRAM, IPPROTO_IP);
    if (sock < 0)
    {
        ESP_LOGE(SSB_TAG, "Failed to create socket. Error %d", errno);
        return -1;
    }
    saddr.sin_family = PF_INET;
    saddr.sin_port = htons(UDP_PORT);
    saddr.sin_addr.s_addr = htonl(INADDR_ANY);
    err = bind(sock, (struct sockaddr *)&saddr, sizeof(struct sockaddr_in));
    if (err < 0)
    {
        ESP_LOGE(SSB_TAG, "Failed to bind socket. Error %d", errno);
        goto err;
    }
    err = socket_add_ipv4_multicast_group(sock, true);
    if (err < 0)
    {
        goto err;
    }
    return sock;
err:
    close(sock);
    return -1;
}

static void udp_group_server_task(void *pvParameters)
{
    char recvbuf[48];
    char echo_str[64]={0};
    extern char device_id[32];
    while (1)
    {
        
        uint32_t bits = 0;
        bits |= IPV4_GOTIP_BIT;
        int sock = create_multicast_ipv4_socket();
        if (sock < 0)
        {
            ESP_LOGE(SSB_TAG, "Failed to create IPv4 multicast socket");
            vTaskDelay(5 / portTICK_PERIOD_MS);
            continue;
        }

        struct sockaddr_in sdestv4 = {
            .sin_family = PF_INET,
            .sin_port = htons(UDP_PORT),
        };
        inet_aton(MULTICAST_IPV4_ADDR, &sdestv4.sin_addr.s_addr);
        int err = 1;
        sprintf(echo_str,"echo:%s",device_id);
        while (err > 0)
        {
            struct timeval tv = {
                .tv_sec = 2,
                .tv_usec = 0,
            };
            fd_set rfds;
            FD_ZERO(&rfds);
            FD_SET(sock, &rfds);
            int s = lwip_select(sock + 1, &rfds, NULL, NULL, &tv);
            if (s < 0)
            {
                ESP_LOGE(SSB_TAG, "Select failed: errno %d", errno);
                err = -1;
                break;
            }
            
            if (s > 0)
            {
                if (FD_ISSET(sock, &rfds))
                {
                    
                    struct sockaddr_in raddr;
                    socklen_t socklen = sizeof(raddr);
                    int len = recvfrom(sock, recvbuf, sizeof(recvbuf) - 1, 0,
                                       (struct sockaddr *)&raddr, &socklen);
                    if (len < 0)
                    {
                        ESP_LOGE(SSB_TAG, "multicast recvfrom failed: errno %d", errno);
                        err = -1;
                        break;
                    }
                    ESP_LOGI(SSB_TAG, "[%s:%hd] -> %d", inet_ntoa(raddr.sin_addr),ntohs(raddr.sin_port), len);
                    
                    sendto(sock, echo_str, strlen(echo_str), 0, (struct sockaddr *)&raddr, socklen);
                }
            }
        }
        ESP_LOGE(SSB_TAG, "Shutting down socket and restarting...");
        shutdown(sock, 0);
        close(sock);
        
    }
}

void SSB_init_lookup_server()
{
    xTaskCreate(udp_group_server_task, "udp_group_server_task", 1024*2, NULL, 10, NULL);
}